#include <linux/kprobes.h>
#include <linux/init.h>
#include <linux/module.h>
#include <linux/sched.h>
#include <linux/slab.h>
#include <linux/unistd.h>
#include <linux/uaccess.h>
#include <linux/fs.h>
#include <linux/namei.h>
#include <linux/ip.h>
#include <linux/etherdevice.h> // 包含以太网头部定义
#include <linux/if_vlan.h>     // 包含 VLAN 相关定义
#include <linux/tcp.h>         // 包含 TCP 头部定义
#include <linux/icmp.h>        // 包含 ICMP 定义
#include <linux/proc_fs.h>     // 包含 procfs 定义

// 定义统计变量
static atomic_t ip_packets = ATOMIC_INIT(0);
static atomic_t tcp_packets = ATOMIC_INIT(0);
static atomic_t udp_packets = ATOMIC_INIT(0);
static atomic_t icmp_packets = ATOMIC_INIT(0);
static atomic_t other_packets = ATOMIC_INIT(0);

// 定义 Kprobe 结构体
static struct kprobe kp = {
    .symbol_name = "ip_rcv", // 监控的函数符号名
};

// 获取当前进程名称
static const char *get_current_comm(void)
{
    struct task_struct *task = current;
    return task->comm;
}

// 定义 pre_handler 回调函数
static int handler_pre(struct kprobe *p, struct pt_regs *regs)
{
    struct sk_buff *skb = (struct sk_buff *)regs->di; // 获取 skb 指针
    struct iphdr *iph = ip_hdr(skb);                  // 获取 IP 头部
    const char *comm = get_current_comm();            // 获取当前进程名称

    // 增加 IP 包计数
    atomic_inc(&ip_packets);

    // 根据 IP 协议类型统计 TCP、UDP 或 ICMP 包
    switch (iph->protocol)
    {
    case IPPROTO_TCP:
        atomic_inc(&tcp_packets);
        break;
    case IPPROTO_UDP:
        atomic_inc(&udp_packets);
        break;
    case IPPROTO_ICMP:
        atomic_inc(&icmp_packets);
        break;
    default:
        atomic_inc(&other_packets);
        break;
    }

    // 打印进程名称和数据包信息
    printk(KERN_INFO "Process: %s, IP Packet received: src=%pI4, dst=%pI4, length=%u, protocol=%d\n",
           comm, &iph->saddr, &iph->daddr, skb->len, iph->protocol);

    return 0;
}

// 定义 post_handler 回调函数
static void handler_post(struct kprobe *p, struct pt_regs *regs, unsigned long flags)
{
    // 可以在这里添加 post_handler 的逻辑
}

// 定义 fault_handler 回调函数
static int handler_fault(struct kprobe *p, struct pt_regs *regs, int trapnr)
{
    printk(KERN_INFO "Kprobe fault at %p, trap #%d\n", p->addr, trapnr);
    return 0;
}

// 定义 proc 文件读取函数
static ssize_t proc_read_stats(struct file *file, char __user *usr_buf, size_t count, loff_t *pos)
{
    char buffer[256];
    int len = 0;

    len += sprintf(buffer + len, "IP Packets: %d\n", atomic_read(&ip_packets));
    len += sprintf(buffer + len, "TCP Packets: %d\n", atomic_read(&tcp_packets));
    len += sprintf(buffer + len, "UDP Packets: %d\n", atomic_read(&udp_packets));
    len += sprintf(buffer + len, "ICMP Packets: %d\n", atomic_read(&icmp_packets));
    len += sprintf(buffer + len, "Other Packets: %d\n", atomic_read(&other_packets));

    if (*pos > 0 || count < len)
        return 0;

    if (copy_to_user(usr_buf, buffer, len))
        return -EFAULT;

    *pos = len;
    return len;
}

// 定义 proc 文件操作结构体
static const struct proc_ops proc_fops = {
    .proc_read = proc_read_stats,
};

// 模块初始化函数
static int __init kprobe_init(void)
{
    // 设置 Kprobe 的回调函数
    kp.pre_handler = handler_pre;
    kp.post_handler = handler_post;
    // kp.fault_handler = handler_fault;

    // 注册 Kprobe
    register_kprobe(&kp);
    printk(KERN_INFO "Kprobe registered\n");

    // 创建 proc 文件
    proc_create("net_stats", 0, NULL, &proc_fops);

    return 0;
}

// 模块卸载函数
static void __exit kprobe_exit(void)
{
    // 注销 Kprobe
    unregister_kprobe(&kp);
    printk(KERN_INFO "Kprobe unregistered\n");

    // 删除 proc 文件
    remove_proc_entry("net_stats", NULL);
}

// 定义模块的初始化和卸载函数
module_init(kprobe_init);
module_exit(kprobe_exit);

MODULE_LICENSE("GPL");
MODULE_DESCRIPTION("Kprobe example to monitor IP packet reception");
MODULE_AUTHOR("Your Name");